package joe;

import java.io.File;
import java.lang.ref.ReferenceQueue;
import java.lang.ref.WeakReference;
import java.util.HashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.function.BiConsumer;
import java.util.function.Function;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import org.lwjgl.system.MemoryUtil;
import org.tinylog.Logger;

public final class ResourceManager {
	private static final class Ref extends WeakReference<Resource> {
		final @NotNull String path;
		final @Nullable Runnable cleaner;

		Ref(@NotNull String path, @NotNull Resource res, @NotNull ReferenceQueue<Resource> refQueue) {
			super(res, refQueue);
			this.path = path;
			cleaner = res.getCleaner();
		}
	}

	public static final @NotNull ExecutorService loaderThreadPool = Executors.newSingleThreadExecutor(ResourceManager::newThread);

	private final ReferenceQueue<Resource> refQueue = new ReferenceQueue<>();
	private final HashMap<@NotNull String, @NotNull String> resDir = new HashMap<>(); // <name,path>
	private final HashMap<@NotNull String, @NotNull Ref> resCache = new HashMap<>(); // <path,>
	private final HashMap<@Nullable String, @NotNull Function<@NotNull String, @Nullable Resource>> loaderMap = new HashMap<>(); // <ext,<path,>>
	private @NotNull BiConsumer<@NotNull String, @NotNull BiConsumer<@NotNull String, @NotNull String>> collector = ResourceManager::defaultCollector; // <path,<name,path>>
	private final HashMap<@NotNull String, Render.@NotNull Shader> vsShaderMap = new HashMap<>();
	private final HashMap<@NotNull String, Render.@NotNull Shader> fsShaderMap = new HashMap<>();
	private final HashMap<@NotNull String, Render.@NotNull Program> programMap = new HashMap<>();
	private final HashMap<@NotNull String, @NotNull Font> fontMap = new HashMap<>();

	ResourceManager() {
		addDefaultLoader();
	}

	void addDefaultLoader() {
		addLoader(null, ResourceMapFile::load);
		for (var ext : new String[]{"jpg", "png", "bmp", "tga", "psd", "pnm", "gif", "avif"})
			addLoader(ext, ResourceTexture::load);
		addLoader("avifs", path -> new ResourceVideo(path, 0, -1));
	}

	void reset() {
		fontMap.clear();
		programMap.clear();
		fsShaderMap.clear();
		vsShaderMap.clear();
		collector = ResourceManager::defaultCollector;
		loaderMap.clear();
		// resCache.clear(); // should be clear by gc & poll
		resDir.clear();
		addDefaultLoader();
	}

	private static @NotNull Thread newThread(@NotNull Runnable r) {
		var thread = new Thread(r, "ResLoader");
		thread.setDaemon(true);
		thread.setPriority(Thread.NORM_PRIORITY - 2);
		return thread;
	}

	public static void defaultCollector(@NotNull String path,
										@NotNull BiConsumer<@NotNull String, @NotNull String> resDirAdder) { // <name,path>
		var file = new File(path);
		if (file.isDirectory()) {
			int n = path.length();
			char c;
			if (n > 0 && ((c = path.charAt(n - 1)) != '/' || c != File.separatorChar))
				path += '/';
			var fs = file.list();
			if (fs != null) {
				for (var f : fs)
					defaultCollector(path + f, resDirAdder);
			}
		} else {
			var name = file.getName();
			resDirAdder.accept(name, path);
			var nameWithoutExt = removePathExt(name);
			if (!nameWithoutExt.equals(name))
				resDirAdder.accept(nameWithoutExt, path);
		}
	}

	public void setCollector(@NotNull BiConsumer<@NotNull String, // <path,>
			@NotNull BiConsumer<@NotNull String, @NotNull String>> collector) { // <name,path>
		this.collector = collector;
	}

	/**
	 * @param loader apply method will run in 'ResLoader' thread
	 */
	public void addLoader(@Nullable String ext, @NotNull Function<@NotNull String, @Nullable Resource> loader) { // <path,>
		loaderMap.put(ext, loader);
	}

	public boolean removeLoader(@Nullable String ext) {
		return loaderMap.remove(ext) != null;
	}

	public void recollectAll(@NotNull String @NotNull ... paths) {
		var t = System.nanoTime();
		resDir.clear();
		BiConsumer<String, String> resDirAdder = resDir::put;
		for (var path : paths)
			collector.accept(path, resDirAdder);
		Logger.debug("ResourceManager: recollectAll collected {} items ({} ms)",
				resDir.size(), (System.nanoTime() - t) / 1_000_000);
	}

	public @Nullable String getPath(@NotNull String name) {
		return resDir.get(name);
	}

	public static @NotNull String getPathExt(@NotNull String path) {
		for (int i = path.length(); --i >= 0; ) {
			var c = path.charAt(i);
			if (c == '.')
				return path.substring(i + 1);
			if (c == '/')
				break;
		}
		return "";
	}

	public static @NotNull String removePathExt(@NotNull String path) {
		for (int i = path.length(); --i >= 0; ) {
			var c = path.charAt(i);
			if (c == '.')
				return path.substring(0, i);
			if (c == '/')
				break;
		}
		return path;
	}

	@SuppressWarnings("unchecked")
	public <R extends Resource> @Nullable R load(@NotNull String name, @NotNull Class<R> resClass) {
		var r = load(name);
		return resClass.isInstance(r) ? (R)r : null;
	}

	public @Nullable Resource load(@NotNull String name) {
		return loadAsync(name, null);
	}

	@SuppressWarnings("unchecked")
	public <R extends Resource> void loadAsync(@NotNull String name, @NotNull Class<R> resClass,
											   @NotNull BiConsumer<@NotNull String, @Nullable R> callback) {
		loadAsync(name, (n, r) -> callback.accept(n, resClass.isInstance(r) ? (R)r : null));
	}

	public @Nullable Resource loadAsync(@NotNull String name,
										@Nullable BiConsumer<@NotNull String, @Nullable Resource> callback) {
		var res = getCache(name);
		if (res != null) {
			if (callback == null)
				return res;
			callback.accept(name, res);
			return null;
		}
		return reload(name, callback);
	}

	public @Nullable Resource getCache(@NotNull String name) {
		var path = getPath(name);
		if (path != null) {
			var wr = resCache.get(path);
			if (wr != null)
				return wr.get();
		}
		return null;
	}

	private @Nullable Resource reload(@NotNull String name,
									  @Nullable BiConsumer<@NotNull String, @Nullable Resource> callback) {
		var path = getPath(name);
		if (path == null) {
			if (callback != null)
				callback.accept(name, null);
			return null;
		}
		var loader = loaderMap.get(getPathExt(path));
		if (loader == null) {
			loader = loaderMap.get(null);
			if (loader == null) {
				if (callback != null)
					callback.accept(name, null);
				return null;
			}
		}
		if (callback == null) {
			var res = loader.apply(path);
			if (res != null)
				resCache.put(path, new Ref(path, res, refQueue));
			return res;
		}
		var loader0 = loader;
		var engine = Engine.current();
		loaderThreadPool.execute(() -> {
			var res = loader0.apply(path);
			engine.postEvent(() -> {
				if (res != null)
					resCache.put(path, new Ref(path, res, refQueue));
				callback.accept(name, res);
			});
		});
		return null;
	}

	public int getResCacheCount() {
		return resCache.size();
	}

	void poll() {
		for (Ref ref; (ref = (Ref)refQueue.poll()) != null; ) {
			resCache.remove(ref.path, ref);
			var cleaner = ref.cleaner;
			if (cleaner != null) {
				try {
					cleaner.run();
				} catch (Exception e) { // logger.error
					Logger.error(e, "ResourceManager: clean exception");
				}
			}
		}
	}

	public @NotNull Render.Program getOrAddShader(@NotNull String vsName, @NotNull String fsName,
												  @NotNull Class<? extends Render.Model> bindModelClass) {
		var progName = vsName + ":" + fsName;
		var prog = programMap.get(progName);
		if (prog != null)
			return prog;
		var engine = Engine.current();
		var render = engine.getRender();
		var vsShader = vsShaderMap.get(vsName);
		if (vsShader == null) {
			var vsPath = engine.getResourceManager().getPath(vsName);
			if (vsPath == null)
				throw new IllegalStateException("not found vsShader: " + vsName);
			var mbb = ResourceMapFile.loadMapFile(vsPath);
			if (mbb == null)
				throw new IllegalStateException("can not load vsShader: " + vsName);
			try {
				vsShader = render.compileShader(Render.Shader.VERTEX, MemoryUtil.memAddress(mbb), mbb.remaining());
			} finally {
				UnsafeUtil.freeMappedByteBuffer(mbb);
			}
			vsShaderMap.put(vsName, vsShader);
		}
		var fsShader = fsShaderMap.get(fsName);
		if (fsShader == null) {
			var fsPath = engine.getResourceManager().getPath(fsName);
			if (fsPath == null)
				throw new IllegalStateException("not found fsShader: " + fsName);
			var mbb = ResourceMapFile.loadMapFile(fsPath);
			if (mbb == null)
				throw new IllegalStateException("can not load fsShader: " + fsName);
			try {
				fsShader = render.compileShader(Render.Shader.FRAGMENT, MemoryUtil.memAddress(mbb), mbb.remaining());
			} finally {
				UnsafeUtil.freeMappedByteBuffer(mbb);
			}
			fsShaderMap.put(fsName, fsShader);
		}
		prog = render.linkShader(vsShader, fsShader, bindModelClass);
		programMap.put(progName, prog);
		return prog;
	}

	public @NotNull Font addFont(@NotNull String fontKey, @NotNull String fontName, int pixelSize, int textureSize) {
		var fontPath = getPath(fontName);
		if (fontPath == null)
			throw new IllegalStateException("not found fontName: " + fontName);
		var font = fontMap.get(fontKey);
		if (font != null) {
			if (font.getPixelSize() == pixelSize && font.getFontPath().equals(fontPath))
				return font;
			fontMap.remove(fontKey);
		}
		font = new Font(Engine.current().getRender(), fontPath, pixelSize, textureSize);
		fontMap.put(fontKey, font);
		return font;
	}

	public @NotNull Font getFont(@NotNull String fontKey) {
		return fontMap.get(fontKey);
	}

	public boolean removeFont(@NotNull String fontKey) {
		return fontMap.remove(fontKey) != null;
	}
}
